package com.atguigu.dga.lineage.service.impl;

import com.atguigu.dga.common.utils.SqlParser;
import com.atguigu.dga.common.utils.SqlUtil;
import com.atguigu.dga.ds.bean.TDsTaskDefinition;
import com.atguigu.dga.ds.service.TDsTaskDefinitionService;
import com.atguigu.dga.lineage.bean.GovernanceLineageTable;
import com.atguigu.dga.lineage.mapper.GovernanceLineageTableMapper;
import com.atguigu.dga.lineage.service.GovernanceLineageTableService;
import com.atguigu.dga.meta.bean.TableMetaInfo;
import com.atguigu.dga.meta.service.TableMetaInfoService;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import lombok.Getter;
import lombok.Setter;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.hive.ql.lib.Dispatcher;
import org.apache.hadoop.hive.ql.lib.Node;
import org.apache.hadoop.hive.ql.parse.ASTNode;
import org.apache.hadoop.hive.ql.parse.HiveParser;
import org.apache.hadoop.hive.ql.parse.SemanticException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import com.baomidou.dynamic.datasource.annotation.DS;

import java.util.*;

/**
 * <p>
 *  服务实现类
 * </p>
 *
 * @author zhangchen
 * @since 2024-08-03
 */
@Service
@DS("dga")
public class GovernanceLineageTableServiceImpl extends ServiceImpl<GovernanceLineageTableMapper, GovernanceLineageTable> implements GovernanceLineageTableService {


    //1  表结构
    //2  获取所有表的元数据信息 (已有)
    //3  提取所有任务的定义 包含了sql
    //4  通过sql 提取其中的 来源表
    //    -> Map<表名，List<来源表>  >
    //         Map<表名，List<输出表>>
    // 保存数据血缘关系表

    @Autowired
    TableMetaInfoService tableMetaInfoService;

    @Autowired
    TDsTaskDefinitionService tDsTaskDefinitionService;

    public void initLineage(String governanceDate) throws Exception {

        // 清理当日已有有的血缘
        remove(new QueryWrapper<GovernanceLineageTable>().eq("governance_date",governanceDate));

        //2  获取所有表的元数据信息 (已有)
        List<TableMetaInfo> tableMetaInfoList = tableMetaInfoService.getListWithExtra(governanceDate);
        //3  提取所有任务的定义 包含了sql
        List<TDsTaskDefinition> taskDefinitionList = tDsTaskDefinitionService.getTaskDefinitionList();

        //4  通过sql 提取其中的 来源表
        //    -> Map<表名，List<来源表>  >
        //         Map<表名，List<输出表>>

        Map<String ,Set<String>> sourceTableNameMap=new HashMap<>();
        Map<String ,Set<String>> sinkTableNameMap=new HashMap<>();
        for (TDsTaskDefinition tDsTaskDefinition : taskDefinitionList) {
            String tableNameWithSchema = tDsTaskDefinition.getName();
            String[] tableNameWithSchemaArr = tableNameWithSchema.split("\\.");
            if(tableNameWithSchemaArr==null||tableNameWithSchemaArr.length!=2){
                continue;
            }
            String schemaName = tableNameWithSchemaArr[0];
            String tableName = tableNameWithSchemaArr[1];

            String sql = tDsTaskDefinition.getSql();
            if(sql==null){
                continue;
            }
            SourceTableDispatcher sourceTableDispatcher = new SourceTableDispatcher();
            sourceTableDispatcher.setDefaultSchemaName(schemaName);
            SqlParser.parseSQL(sql,sourceTableDispatcher);
            //提取成果
            Set<String> sourceTableSet = sourceTableDispatcher.getSourceTableSet();
            //过滤掉自己
            sourceTableSet.remove(tableNameWithSchema);

            //装配来源表   m->Set( a,b,c)    n->Set(a,d)
            sourceTableNameMap.put(tableNameWithSchema,sourceTableSet);
            //装配输出表   1 a-> Set(m)  b->Set(m) c->Set(m)  2  a->Set(m,n)  b->Set(m) c->Set(m) d->Set(n)
            for (String sourceTableName : sourceTableSet) {
                //先检查map是否已经有该键值
                Set<String> sinkTableSet = sinkTableNameMap.get(sourceTableName);
                if(sinkTableSet==null){
                    sinkTableSet = new HashSet<>();
                }
                sinkTableSet.add(tableNameWithSchema);
                sinkTableNameMap.put(sourceTableName,sinkTableSet);
            }
        }
        // 装配数据血缘关系表
        List<GovernanceLineageTable> lineageTableList=new ArrayList<>();
        for (TableMetaInfo tableMetaInfo : tableMetaInfoList) {
            GovernanceLineageTable governanceLineageTable = new GovernanceLineageTable();
            governanceLineageTable.setGovernanceDate(governanceDate);
            governanceLineageTable.setSchemaName(tableMetaInfo.getSchemaName());
            governanceLineageTable.setTableName(tableMetaInfo.getTableName());
            governanceLineageTable.setCreateTime(new Date());

            Set<String> sourceTableSet = sourceTableNameMap.get(tableMetaInfo.getSchemaName() + "." + tableMetaInfo.getTableName());
            governanceLineageTable.setSourceTables(StringUtils.join(sourceTableSet,","));

            Set<String> sinkTableSet = sinkTableNameMap.get(tableMetaInfo.getSchemaName() + "." + tableMetaInfo.getTableName());
            governanceLineageTable.setSinkTables(StringUtils.join(sinkTableSet,","));

            lineageTableList.add(governanceLineageTable);

        }
        //保存
        saveBatch(lineageTableList);

    }

    class SourceTableDispatcher implements Dispatcher {

        @Getter
        Set<String> sourceTableSet=new HashSet<>();

        @Setter
        String defaultSchemaName;


        @Override
        public Object dispatch(Node nd, Stack<Node> stack, Object... nodeOutputs) throws SemanticException {
            ASTNode astNode = (ASTNode) nd;
            if(astNode.getType()== HiveParser.TOK_TABREF){
                if(astNode.getChild(0).getChildCount()>1){ //有库名
                    String schemaName = astNode.getChild(0).getChild(0).getText();
                    String tableName = astNode.getChild(0).getChild(1).getText();
                    sourceTableSet.add(schemaName+"."+tableName);

                }else{//无库名
                    String tableName = astNode.getChild(0).getChild(0).getText();
                    sourceTableSet.add(defaultSchemaName+"."+tableName);
                }

            }


            return null;
        }
    }


}
